# -*- coding: UTF-8 -*-
# @IDE     : VScode
# @File   : chatBot_streamlit_app.py
# @Time   : 2024/05/26 17:24:55
# @Author : zhonggc

import streamlit as st
from dotenv import load_dotenv, find_dotenv
from format_chain import chat_history, get_answer
from st_chat_message import message as chat_message
from streamlit_option_menu import option_menu
from chat_history import ChatHistory
from uuid import uuid1

# TODO 1. 侧边会话选择功能 2. 聊天记录保存到数据库


chat_history_hander = ChatHistory()


class StreamlitApp:

    session_id = "test"

    def __init__(self):
        self.setup()

    def setup(self):
        _ = load_dotenv(find_dotenv())
        self.init_st_layout()
        self.load_chat_history()
        self.get_answer()

    def init_st_layout(self):
        st.title("🤡AI知识助手")
        # TODO 侧边会话选择功能无法将会话title和会话id对应
        session_list = chat_history_hander.pretty_session_dict()
        with st.sidebar:
            if not session_list:
                st.warning("暂无会话")
            else:
                self.session_id = option_menu(
                    "会话选择",
                    [session["session_id"] for session in session_list],
                    default_index=0,
                    orientation="vertical",
                    key="session_id",
                )
            if st.button(label="新建会话", key="new_session"):
                self.create_new_session()
        self.prompt = st.chat_input('给"AI助手"发送消息')

    def create_new_session(self):
        """新建会话"""
        chat_history.messages.clear()
        chat_history_hander.save_chat_history(messages=chat_history.messages)
        # 刷新页面
        st.rerun()

    def load_chat_history(self):
        # 刷新页面时要重新加载内存中的聊天记录
        chat_history.messages = chat_history_hander.load_chat_history(self.session_id)
        for index, message in enumerate(chat_history.messages):
            if message.type == "human":
                chat_message(message.content, is_user=True, key=index)
            else:
                chat_message(message.content, is_user=False, key=index)

    def get_answer(self):
        """获取答案"""
        if self.prompt:
            chat_message(self.prompt, is_user=True, key=str(uuid1()))
            response = get_answer(self.prompt, self.session_id)
            if response:
                chat_message(response, is_user=False, key=str(uuid1()))

            # 保存聊天记录到数据库
            chat_history_hander.save_chat_history(
                self.session_id, chat_history.messages
            )


if __name__ == "__main__":
    app = StreamlitApp()
